/*=========================================================================================
  This file is a part of BioDao.
  Copyright (C) 2011, Fu Limin. Email: fu@daovm.net, limin.fu@yahoo.com

  This software is free software; you can redistribute it and/or modify it under the terms
  of the GNU Lesser General Public License as published by the Free Software Foundation;
  either version 2.1 of the License, or (at your option) any later version.

  This software is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
  without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
  See the GNU Lesser General Public License for more details.
  =========================================================================================*/

#include<daoStream.h>
#include"daoxSequence.h"
#include"daoxFastAQ.h"

DAO_INIT_MODULE

static DaoValue* FastAQ_MakeValue( int type, DString *ds, DString *ss, DString *qs, DaoValue *value )
{
	switch( type ){
	case DAO_FASTAQ_DESC : return value;
	case DAO_FASTAQ_AUTO : return value;
	case DAO_FASTAQ_FASTA : return (DaoValue*) DaoxSequence_New( ds, ss );
	case DAO_FASTAQ_FASTQ : return (DaoValue*) DaoxSequenceQS_New( ds, ss, qs );
	}
	return value;
}

int DaoxFastAQ_Read( FILE *fin, DArray *seqs, int type, int count )
{
	int getdes = 1;
	DString *line, *ds, *ss, *qs, *data=NULL;
	DaoString desc = {DAO_STRING,0,0,0,1,NULL};
	DaoString *desc2 = & desc;
	DaoValue *value = (DaoValue*) desc2;
	DaoxSequenceQS *seq;

	if( fin == NULL ) return 0;
	line = DString_New(1);
	seq = DaoxSequenceQS_New( line, line, line );
	ds = seq->des;
	ss = seq->seq;
	qs = seq->qs;
	desc.data = ds;
	while( feof( fin ) == 0 ){
		long pos = ftell( fin );
		if( DaoFile_ReadLine( fin, line ) == 0 ) break;
		if( getdes && (line->mbs[0] == '>' || line->mbs[0] == '@' || line->mbs[0] == '+') ){
			getdes = 0;
			if( seqs->size == 0 && type == DAO_FASTAQ_AUTO ){
				type = line->mbs[0] == '>' ? DAO_FASTAQ_FASTA : DAO_FASTAQ_FASTQ;
			}
			DString_Chop( line );
			if( line->mbs[0] == '+' ){
				data = qs;
			}else{
				if( ds->size || ss->size ){
					DArray_Append( seqs, FastAQ_MakeValue( type, ds, ss, qs, value ) );
					DString_Reset( ds, 0 );
					DString_Reset( ss, 0 );
					DString_Reset( qs, 0 );
					if( count && seqs->size >= count ){
						fseek( fin, pos, SEEK_SET );
						break;
					}
				}
				data = ss;
				DString_SetMBS( ds, line->mbs + 1 );
				if( type == DAO_FASTAQ_DESC ){
					int d = 0;
					while( d < ds->size && isspace( ds->mbs[d] ) ==0 ) d += 1;
					DString_Resize( ds, d );
				}
			}
		}else{
			DString_AppendMBS( data, line->mbs );
			if( data->size && isspace( data->mbs[data->size-1] ) ){
				data->size -= 1;
				data->mbs[ data->size ] = 0;
			}
			getdes = data == ss || (data == qs && qs->size == ss->size);
		}
	}
	if( ds->size || ss->size ) DArray_Append( seqs, FastAQ_MakeValue( type, ds, ss, qs, value ) );
	DString_Delete( line );
	DaoxSequenceQS_Delete( seq );
	return 1;
}

static int DaoFastAQ_Convert( FILE *fin, FILE *fout )
{
	int getdes = 1;
	DString *line, *ss, *qs, *data=NULL;

	if( fin == NULL ) return 0;
	ss = DString_New(1);
	qs = DString_New(1);
	line = DString_New(1);
	while( feof( fin ) == 0 ){
		if( DaoFile_ReadLine( fin, line ) == 0 ) break;
		if( getdes && (line->mbs[0] == '>' || line->mbs[0] == '@' || line->mbs[0] == '+') ){
			getdes = 0;
			if( line->mbs[0] == '+' ){
				data = qs;
			}else{
				data = ss;
				DString_Reset( ss, 0 );
				DString_Reset( qs, 0 );
				fprintf( fout, ">%s", line->mbs + 1 );
			}
		}else{
			if( data == ss ) fprintf( fout, "%s", line->mbs );
			getdes = data == ss || (data == qs && qs->size == ss->size);
		}
	}
	DString_Delete( ss );
	DString_Delete( qs );
	DString_Delete( line );
	return 1;
}
void DaoxSequence_WriteFastAQ( DaoxSequence *self, FILE *fout, int W )
{
	DString *ds = self->des;
	DString *ss = self->seq;
	int i;
	fprintf( fout, ">%s\n", ds->mbs );
	if( W ){
		for(i=0; i<ss->size; i++){
			fprintf( fout, "%c", ss->mbs[i] );
			if( (i+1) % W ==0 || (i+1) == ss->size ) fprintf( fout, "\n" );
		}
	}else{
		fprintf( fout, "%s\n", ss->mbs );
	}
}
void DaoxSequenceQS_WriteFastAQ( DaoxSequenceQS *self, FILE *fout, int W )
{
	DString *ds = self->des;
	DString *ss = self->seq;
	DString *qs = self->qs;
	int i;
	fprintf( fout, "@%s\n", ds->mbs );
	if( W ){
		for(i=0; i<ss->size; i++){
			fprintf( fout, "%c", ss->mbs[i] );
			if( (i+1) % W ==0 || (i+1) == ss->size ) fprintf( fout, "\n" );
		}
	}else{
		fprintf( fout, "%s\n", ss->mbs );
	}
	fprintf( fout, "+%s\n", ds->mbs );
	if( W ){
		for(i=0; i<qs->size; i++){
			fprintf( fout, "%c", qs->mbs[i] );
			if( (i+1) % W ==0 || (i+1) == qs->size ) fprintf( fout, "\n" );
		}
	}else{
		fprintf( fout, "%s\n", qs->mbs );
	}
}
void DaoxFastAQ_Write( FILE *fout, DArray *seqs, int type, int width )
{
	size_t i;
	for(i=0; i<seqs->size; i++){
		DaoxSequence *seq = (DaoxSequence*) seqs->items.pValue[i];
		if( type == DAO_FASTAQ_FASTA ){
			DaoxSequence_WriteFastAQ( seq, fout, width );
		}else if( seq->typer == & DaoxSequenceQS_Typer ){
			DaoxSequenceQS_WriteFastAQ( (DaoxSequenceQS*) seq, fout, width );
		}else{
			DaoxSequence_WriteFastAQ( seq, fout, width );
		}
	}
}

static void FastAQ_ReadFile( DaoProcess *proc, DString *input, int type )
{
	DaoList *list = DaoProcess_PutList( proc );
	FILE *fin = fopen( DString_GetMBS( input ), "r" );
	if( fin == NULL ){
		DaoProcess_RaiseException( proc, DAO_ERROR, "failed to open file" );
		return;
	}
	DaoxFastAQ_Read( fin, & list->items, type, 0 );
}
static void FastAQ_ReadStream( DaoProcess *proc, DaoStream *input, int type, int count )
{
	DaoList *list = DaoProcess_PutList( proc );
	FILE *fin = stdin;
	if( input->file && input->file->fd ) fin = input->file->fd;
	if( (input->mode & DAO_IO_READ) == 0 ){
		DaoProcess_RaiseException( proc, DAO_ERROR, "stream is not readable" );
		return;
	}
	DaoxFastAQ_Read( fin, & list->items, type, count );
}
static void FastAQ_ReadF( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadFile( proc, p[0]->xString.data, DAO_FASTAQ_AUTO );
}
static void FastAQ_ReadFA( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadFile( proc, p[0]->xString.data, DAO_FASTAQ_FASTA );
}
static void FastAQ_ReadFQ( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadFile( proc, p[0]->xString.data, DAO_FASTAQ_FASTQ );
}
static void FastAQ_ReadFD( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadFile( proc, p[0]->xString.data, DAO_FASTAQ_DESC );
}
static void FastAQ_ReadS( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadStream( proc, (DaoStream*) p[0], DAO_FASTAQ_AUTO, p[1]->xInteger.value );
}
static void FastAQ_ReadSA( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadStream( proc, (DaoStream*) p[0], DAO_FASTAQ_FASTA, p[2]->xInteger.value );
}
static void FastAQ_ReadSQ( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadStream( proc, (DaoStream*) p[0], DAO_FASTAQ_FASTQ, p[2]->xInteger.value );
}
static void FastAQ_ReadSD( DaoProcess *proc, DaoValue *p[], int N )
{
	FastAQ_ReadStream( proc, (DaoStream*) p[0], DAO_FASTAQ_DESC, p[2]->xInteger.value );
}
static void FastAQ_Writes( DaoProcess *proc, DaoValue *p[], int N )
{
	DArray *seqs = & p[1]->xList.items;
	int width  = p[2]->xInteger.value;
	int type = p[3]->xInteger.value == 0 ? DAO_FASTAQ_FASTA : DAO_FASTAQ_AUTO;
	FILE *fout = stdout;
	if( p[0]->type == DAO_STREAM ){
		DaoStream *stream = (DaoStream*) p[0];
		if( stream->file && stream->file->fd ) fout = stream->file->fd;
		if( (stream->mode & DAO_IO_WRITE) == 0 ){
			DaoProcess_RaiseException( proc, DAO_ERROR, "stream is not writable" );
			return;
		}
		DaoxFastAQ_Write( fout, seqs, type, width );
	}else{
		fout = fopen( p[0]->xString.data->mbs, "w+" );
		if( fout == NULL ){
			DaoProcess_RaiseException( proc, DAO_ERROR, "failed to open file" );
			return;
		}
		DaoxFastAQ_Write( fout, seqs, type, width );
		fclose( fout );
	}
}
static void FastAQ_Write( DaoProcess *proc, DaoValue *p[], int N )
{
	DaoStream *stream = (DaoStream*) p[0];
	DaoxSequence *seq = (DaoxSequence*) p[1];
	int width = p[2]->xInteger.value;
	FILE *fout = stdout;
	if( stream->file && stream->file->fd ) fout = stream->file->fd;
	if( p[3]->xInteger.value == 0 ){
		DaoxSequence_WriteFastAQ( seq, fout, width );
	}else if( seq->typer == & DaoxSequenceQS_Typer ){
		DaoxSequenceQS_WriteFastAQ( (DaoxSequenceQS*) seq, fout, width );
	}else{
		DaoxSequence_WriteFastAQ( seq, fout, width );
	}
}
static void FastAQ_Convert( DaoProcess *proc, DaoValue *p[], int N )
{
	FILE *fin, *fout;
	DString *fastq = p[0]->xString.data;
	DString *fasta = p[1]->xString.data;
	DString_ToMBS( fastq );
	DString_ToMBS( fasta );
	fin = fopen( fastq->mbs, "r" );
	fout = fopen( fasta->mbs, "w+" );
	if( fin == NULL || fout == NULL ){
		DaoProcess_RaiseException( proc, DAO_ERROR, "failed to open file" );
		return;
	}
	DaoFastAQ_Convert( fin, fout );
	fclose( fin );
	fclose( fout );
}

static DaoFuncItem DaoFastAQ_Meths[]=
{
	{ FastAQ_ReadF,   "Read( fin :string ) => list<Sequence>" },
	{ FastAQ_ReadFA,  "Read( fin :string, format :enum<fasta> ) => list<Sequence>" },
	{ FastAQ_ReadFQ,  "Read( fin :string, format :enum<fastq> ) => list<SequenceQS>" },
	{ FastAQ_ReadFD,  "Read( fin :string, desc :enum<description> ) => list<string>" },
	{ FastAQ_ReadS,   "Read( fin :stream, count = 0 ) => list<Sequence>" },
	{ FastAQ_ReadSA,  "Read( fin :stream, format :enum<fasta>, count = 0 ) => list<Sequence>" },
	{ FastAQ_ReadSQ,  "Read( fin :stream, format :enum<fastq>, count = 0 ) => list<SequenceQS>" },
	{ FastAQ_ReadSD,  "Read( fin :stream, desc :enum<description>, count = 0 ) => list<string>" },

	{ FastAQ_Writes,  "Write( fout :string, seq :list<Sequence>, width = 80, qscore = 1 )" },
	{ FastAQ_Writes,  "Write( fout :stream, seq :list<Sequence>, width = 80, qscore = 1 )" },
	{ FastAQ_Write,   "Write( fout :stream, seq :Sequence, width = 80, qscore = 1 )" },

	{ FastAQ_Convert, "Convert( fin :string, fout :string )" },
	{ NULL, NULL }
};

static DaoTypeBase DaoFastAQ_Typer = { "FastAQ", NULL, NULL, DaoFastAQ_Meths, {0}, {0}, NULL, NULL };


int DaoOnLoad( DaoVmSpace *vmSpace, DaoNamespace *ns )
{
	DaoVmSpace_LinkModule( vmSpace, ns, "bio/core/type/sequence" );
	DaoNamespace_WrapType( ns, & DaoFastAQ_Typer, 1 );
	return 0;
}
